# %%
from zhipuai import ZhipuAI
client = ZhipuAI(api_key="") # 请填写您自己的APIKey
# %%
# case_01 = """
#     (a) Married individuals filing joint returns and surviving spouses
#     There is hereby imposed on the taxable income of-
#     (1) every married individual (as defined in section 7703) who makes a single return jointly with his spouse under section 6013, and
#     (2) every surviving spouse (as defined in section 2(a)),
#     a tax determined in accordance with the following table:

#     If taxable income is: 	The tax is:
#     Not over $36,900 	15% of taxable income.
#     Over $36,900 but not over $89,150 	$5,535, plus 28% of the excess over $36,900.
#     Over $89,150 but not over $140,000 	$20,165, plus 31% of the excess over $89,150.
#     Over $140,000 but not over $250,000 	$35,928.50, plus 36% of the excess over $140,000.
#     Over $250,000 	$75,528.50, plus 39.6% of the excess over $250,000.
#     """

with open("case01_in.txt", "r", encoding="utf-8") as f:
    case01_in = f.read()

prompts = []
prompts.append(case01_in)
prompts.append("请帮我介绍一下这段文字的含义")
prompts.append("请对比中国的个人所得税法指出这段规定的特殊性")
prompts.append("根据这段文字描述给出一个python函数用于计算这种情况下的美国个人所得税")
prompts.append("我想继续了解关于这种情况的美国个人所得税相关计算与规定，可以以什么关键词在税法中搜索？")

contents = []
messages = []
messages.append({"role": "system", "content": "美国税制专家"})

for prompt in prompts:
    contents.append(prompt)
    messages.append({"role": "user", "content": prompt})
    response = client.chat.completions.create(
        model="GLM-4-Flash-250414",  # 填写需要调用的模型编码
        messages = messages,
    )
    content = response.choices[0].message.content
    messages.append({"role": "assistant", "content": content})
    contents.append(content)

with open("case01_out.txt", "w", encoding="utf-8") as f:
    f.write("\n".join(contents))

# %%

def calculate_us_income_tax(income):
    """
    Calculate the US income tax based on the provided brackets for married filing jointly or surviving spouses.

    Args:
        income (float): The taxable income of the individual.

    Returns:
        float: The calculated income tax.
    """
    if income <= 36900:
        tax = 0.15 * income
    elif income <= 89150:
        tax = 5535 + 0.28 * (income - 36900)
    elif income <= 140000:
        tax = 20165 + 0.31 * (income - 89150)
    elif income <= 250000:
        tax = 35928.50 + 0.36 * (income - 140000)
    else:
        tax = 75528.50 + 0.396 * (income - 250000)

    return tax

calculate_us_income_tax(1e5)

# %%
